﻿using JESAI.DistributedLock.Internal;
using JESAI.DistributedLock.Internal.Data;
using Npgsql;
using System;
using System.Data;
using System.Text;
using System.Text.RegularExpressions;
using System.Threading;
using System.Threading.Tasks;

namespace JESAI.DistributedLock.Postgres;

/// <summary>
/// 实现 <see cref="IDbSynchronizationStrategy{TLockCookie}"/> using advisory locking functions
/// (see https://www.postgresql.org/docs/12/functions-admin.html#FUNCTIONS-ADVISORY-LOCKS)
/// </summary>
internal class PostgresAdvisoryLock : IDbSynchronizationStrategy<object>
{
    private static readonly object Cookie = new();

    public static readonly PostgresAdvisoryLock ExclusiveLock = new(isShared: false),
        SharedLock = new(isShared: true);

    private readonly bool _isShared;

    private PostgresAdvisoryLock(bool isShared)
    {
        this._isShared = isShared;
    }

    /// <summary>
    /// 咨询锁本身不支持可升级
    /// </summary>
    public bool IsUpgradeable => false;

    public async ValueTask<object?> TryAcquireAsync(DatabaseConnection connection, string resourceName, TimeoutValue timeout, CancellationToken cancellationToken)
    {
        const string SavePointName = "medallion_threading_postgres_advisory_lock_acquire";

        PostgresAdvisoryLockKey key = new(resourceName);

        if (connection.IsExernallyOwned
            && await this.IsHoldingLockAsync(connection, key, cancellationToken).ConfigureAwait(false))
        {
            if (timeout.IsZero) { return null; }
            if (timeout.IsInfinite) { throw new DeadlockException("Attempted to acquire a lock that is already held on the same connection"); }
            await SyncViaAsync.Delay(timeout, cancellationToken).ConfigureAwait(false);
            return null;
        }

        // 只有在我们尝试获取事务范围的锁的情况下，我们才会定义一个保存点，但如果成功获取锁，我们将无法将其回滚，因为锁将被释放。因此，在这种情况下，我们在设置保存点之前捕获超时设置值，然后在尝试获取锁后尝试恢复这些值。
        // 注意：为了捕获和恢复所有情况下的值，不能删除保存点功能。
        // 当尝试获取锁时发生错误时，事务将中止，我们无法运行任何其他查询，除非事务或保存点被回滚。
        var capturedTimeoutSettings = await CaptureTimeoutSettingsIfNeededAsync(connection, cancellationToken).ConfigureAwait(false);

        // 我们的acquire命令将使用SET LOCAL设置语句超时。如果我们处于事务中，这将持续到当前事务结束，而不仅仅是当前批处理结束。为了确保我们不会泄露这些设置，在事务的情况下，我们首先设置了一个保存点，稍后可以回滚（只有在我们没有获得事务范围的锁的情况下——使用设置更改，但不使用锁）。因为在不知道保存点已设置的情况下，我们无法自信地回滚它，所以在尝试catch之前，我们会在自己的查询中启动保存点。
        var needsSavePoint = await ShouldDefineSavePoint(connection).ConfigureAwait(false);

        if (needsSavePoint)
        {
            using var setSavePointCommand = connection.CreateCommand();
            setSavePointCommand.SetCommandText("SAVEPOINT " + SavePointName);
            await setSavePointCommand.ExecuteNonQueryAsync(cancellationToken).ConfigureAwait(false);
        }

        using var acquireCommand = this.CreateAcquireCommand(connection, key, timeout);

        object? acquireCommandResult;
        try
        {
            acquireCommandResult = await acquireCommand.ExecuteScalarAsync(cancellationToken).ConfigureAwait(false);
        }
        catch (Exception ex)
        {
            await RollBackTransactionTimeoutVariablesIfNeededAsync(acquired: false).ConfigureAwait(false);

            await RestoreTimeoutSettingsIfNeededAsync(capturedTimeoutSettings, connection).ConfigureAwait(false);

            if (ex is PostgresException postgresException)
            {
                switch (postgresException.SqlState)
                {
                    // 来自的lock_timeout错误代码 https://www.postgresql.org/docs/16/errcodes-appendix.html
                    case "55P03":
                        // 即使我们遇到了锁超时，Postgres中的一个潜在竞争条件意味着我们实际上可能在超时之前就获得了锁。为了解释这一点，我们只需重新检查是否持有锁以确定最终返回值。请参阅https://github.com/madelson/DistributedLock/issues/147以及https://www.postgresql.org/message-id/63573.1668271677%40sss.pgh.pa.us了解更多详情。
                        //注意：我们使用CancellationToken。此检查无结果，因为如果我们持有锁，则中止将无效。
                        return await this.IsHoldingLockAsync(connection, key, CancellationToken.None).ConfigureAwait(false)
                            ? Cookie
                            : null;
                    //死锁检测到的错误代码来自 https://www.postgresql.org/docs/16/errcodes-appendix.html
                    case "40P01":
                        throw new DeadlockException($"The request for the distributed lock failed with exit code '{postgresException.SqlState}' (deadlock_detected)", ex);
                }
            }

            if (ex is OperationCanceledException
                && cancellationToken.IsCancellationRequested
                // 除了回滚之外，没有其他方法可以显式释放事务范围的锁；就我们而言
                // 回滚事务超时变量IfNeededAsync将已通过回滚保存点释放。
                // 此外，调用者将继续处理交易。
                && !UseTransactionScopedLock(connection))
            {
                // 如果我们在收购过程中退出，请确保我们没有留下锁
                await this.ReleaseAsync(connection, key, isTry: true).ConfigureAwait(false);
            }

            throw;
        }

        var acquired = acquireCommandResult switch
        {
            DBNull _ => true, // 表示我们调用了pg_advisory_lock，而不是pg_try_advisory.lock
            null => true, // Npgsql 8返回null而不是DBNull
            false => false,
            true => true,
            _ => default(bool?)
        };

        await RollBackTransactionTimeoutVariablesIfNeededAsync(acquired: acquired == true).ConfigureAwait(false);

        await RestoreTimeoutSettingsIfNeededAsync(capturedTimeoutSettings, connection).ConfigureAwait(false);

        return acquired switch
        {
            false => null,
            true => Cookie,
            null => throw new InvalidOperationException($"Unexpected value '{acquireCommandResult}' from acquire command")
        };

        async ValueTask RollBackTransactionTimeoutVariablesIfNeededAsync(bool acquired)
        {
            if (needsSavePoint
                // 对于事务范围的锁，我们不能在成功时回滚保存点，因为这将回滚我们对锁的持有。
                // “泄漏”保存点是可以的，因为如果它是内部拥有的事务，那么保存点将随着事务的处置而被清理。
                // 如果这是一笔外部拥有的交易，那么我们必须“泄露”它，否则我们将失去锁。此外，在这种情况下，我们无法避免使用保存点，因为否则如果发生异常，外部拥有的事务将被中止并完全不可用。
                && !(acquired && UseTransactionScopedLock(connection)))
            {
                // 尝试清除我们设置的超时变量
                using var rollBackSavePointCommand = connection.CreateCommand();
                rollBackSavePointCommand.SetCommandText("ROLLBACK TO SAVEPOINT " + SavePointName);
                await rollBackSavePointCommand.ExecuteNonQueryAsync(CancellationToken.None).ConfigureAwait(false);
            }
        }
    }

    private async Task<bool> IsHoldingLockAsync(DatabaseConnection connection, PostgresAdvisoryLockKey key, CancellationToken cancellationToken)
    {
        using var command = connection.CreateCommand();
        command.SetCommandText($@"
                SELECT COUNT(*) 
                FROM pg_catalog.pg_locks l
                JOIN pg_catalog.pg_database d
                    ON d.oid = l.database
                WHERE l.locktype = 'advisory' 
                    AND {AddPGLocksFilterParametersAndGetFilterExpression(command, key)} 
                    AND l.pid = pg_catalog.pg_backend_pid() 
                    AND d.datname = pg_catalog.current_database()"
        );
        return (long)(await command.ExecuteScalarAsync(cancellationToken).ConfigureAwait(false))! != 0;
    }

    private DatabaseCommand CreateAcquireCommand(DatabaseConnection connection, PostgresAdvisoryLockKey key, TimeoutValue timeout)
    {
        var command = connection.CreateCommand();

        var commandText = new StringBuilder();

        // 我们将statement_timeout设置为0（inf），因为我们希望所有内容都由lock_timeout驱动。
        commandText.AppendLine("SET LOCAL statement_timeout = 0;");
        // 我们将锁超时设置为超时，但如果超时为零，则将其设置为inf，因为我们将使用pg_try_advisory_lock函数，在这种情况下不会阻塞。
        commandText.AppendLine($"SET LOCAL lock_timeout = {(timeout.IsZero || timeout.IsInfinite ? 0 : timeout.InMilliseconds)};");

        commandText.Append("SELECT ");
        var isTry = timeout.IsZero;
        commandText.Append("pg_catalog.pg");
        if (isTry) { commandText.Append("_try"); }
        commandText.Append("_advisory");
        if (UseTransactionScopedLock(connection)) { commandText.Append("_xact"); }
        commandText.Append("_lock");
        if (this._isShared) { commandText.Append("_shared"); }
        commandText.Append('(').Append(AddKeyParametersAndGetKeyArguments(command, key)).Append(')')
            .Append(" AS result");

        command.SetCommandText(commandText.ToString());
        command.SetTimeout(timeout);

        return command;
    }

    private static async ValueTask<CapturedTimeoutSettings?> CaptureTimeoutSettingsIfNeededAsync(DatabaseConnection connection, CancellationToken cancellationToken)
    {
        var shouldCaptureTimeoutSettings = connection.IsExernallyOwned && UseTransactionScopedLock(connection);

        // 如果我们不尝试获取外部拥有的事务范围锁，则返回null。
        if (!shouldCaptureTimeoutSettings) { return null; }

        var statementTimeout = await GetCurrentSetting("statement_timeout", connection, cancellationToken).ConfigureAwait(false);
        var lockTimeout = await GetCurrentSetting("lock_timeout", connection, cancellationToken).ConfigureAwait(false);

        var capturedTimeoutSettings = new CapturedTimeoutSettings(statementTimeout!, lockTimeout!);

        return capturedTimeoutSettings;

        async ValueTask<string?> GetCurrentSetting(string settingName, DatabaseConnection connection, CancellationToken cancellationToken)
        {
            using var getCurrentSettingCommand = connection.CreateCommand();

            getCurrentSettingCommand.SetCommandText($"SELECT current_setting('{settingName}', 'true') AS {settingName};");

            return (string?)await getCurrentSettingCommand.ExecuteScalarAsync(cancellationToken).ConfigureAwait(false);
        }
    }

    private static async ValueTask<bool> ShouldDefineSavePoint(DatabaseConnection connection)
    {
        // 如果连接是内部拥有的，我们只在事务已打开时定义一个保存点。
        if (!connection.IsExernallyOwned) { return connection.HasTransaction; }

        // 如果连接由已建立事务的外部拥有，则意味着连接是通过事务锁定API实现的（请参阅PostgresDistributeLock.Transactions.cs）。
        if (connection.HasTransaction) { return true; }

        // 外部拥有的连接可能仍然是我们看不到的交易的一部分。
        // 只有当外部拥有的连接没有通过事务锁定API时才会出现这种情况（请参阅PostgresDistributeLock.Transactions.cs）。
        // 在这种情况下，检测交易的唯一真正方法是开始一个新的交易。
        try
        {
            await connection.BeginTransactionAsync().ConfigureAwait(false);
        }
        catch (InvalidOperationException)
        {
            // 与事务的外部拥有连接=>我们需要定义一个保存点。
            return true;
        }

        await connection.DisposeTransactionAsync().ConfigureAwait(false);

        // 无交易的外部拥有连接=>无保存点
        return false;
    }

    private static async ValueTask RestoreTimeoutSettingsIfNeededAsync(CapturedTimeoutSettings? settings, DatabaseConnection connection)
    {
        // 如果我们没有尝试获取外部拥有的事务范围锁，则设置应为null。
        if (settings is null) { return; }

        using var restoreTimeoutSettingsCommand = connection.CreateCommand();

        StringBuilder commandText = new();
        commandText.AppendLine($"SET LOCAL statement_timeout = {settings.Value.StatementTimeout};");
        commandText.AppendLine($"SET LOCAL lock_timeout = {settings.Value.LockTimeout};");

        restoreTimeoutSettingsCommand.SetCommandText(commandText.ToString());

        await restoreTimeoutSettingsCommand.ExecuteNonQueryAsync(CancellationToken.None).ConfigureAwait(false);
    }

    public ValueTask ReleaseAsync(DatabaseConnection connection, string resourceName, object lockCookie) =>
        this.ReleaseAsync(connection, new PostgresAdvisoryLockKey(resourceName), isTry: false);

    private async ValueTask ReleaseAsync(DatabaseConnection connection, PostgresAdvisoryLockKey key, bool isTry)
    {
        Invariant.Require(!UseTransactionScopedLock(connection));

        using var command = connection.CreateCommand();
        command.SetCommandText($"SELECT pg_catalog.pg_advisory_unlock{(this._isShared ? "_shared" : string.Empty)}({AddKeyParametersAndGetKeyArguments(command, key)})");
        var result = (bool)(await command.ExecuteScalarAsync(CancellationToken.None).ConfigureAwait(false))!;
        if (!isTry && !result)
        {
            throw new InvalidOperationException("Attempted to release a lock that was not held");
        }
    }

    private static string AddKeyParametersAndGetKeyArguments(DatabaseCommand command, PostgresAdvisoryLockKey key)
    {
        if (key.HasSingleKey)
        {
            command.AddParameter("key", key.Key, DbType.Int64);
            return "@key";
        }
        else
        {
            var (key1, key2) = key.Keys;
            command.AddParameter("key1", key1, DbType.Int32);
            command.AddParameter("key2", key2, DbType.Int32);
            return "@key1, @key2";
        }
    }

    private static bool UseTransactionScopedLock(DatabaseConnection connection) =>
        // 事务范围的锁定支持内部拥有的连接和明确具有事务的外部拥有的连接（这意味着外部连接是通过事务锁定API实现的，请参阅PostgresDistributeLock.Transactions.cs）。
        connection.HasTransaction;

    private static string AddPGLocksFilterParametersAndGetFilterExpression(DatabaseCommand command, PostgresAdvisoryLockKey key)
    {
        // 来自 https://www.postgresql.org/docs/12/view-pg-locks.html
        // 可以在由单个bigint值或两个整数值组成的键上获取咨询锁。
        // bigint键的高阶半部分显示在classid列中，低阶半部分在objid列中并且objidd等于1。原始的bigint值可以用表达式（classd:：bigint<<32）|objid:：bigint重新组装。
        // 整数键显示为第一个键在classid列，第二个键在objid列，objsubid等于2。

        string classIdParameter, objIdParameter, objSubId;
        if (key.HasSingleKey)
        {
            // 由于Postgres似乎缺少未经检查的int转换，在这种情况下，只生成额外的参数来携带拆分键信息会更简单
            var (keyUpper32, keyLower32) = key.Keys;
            command.AddParameter(classIdParameter = "keyUpper32", keyUpper32, DbType.Int32);
            command.AddParameter(objIdParameter = "keyLower32", keyLower32, DbType.Int32);
            objSubId = "1";
        }
        else
        {
            AddKeyParametersAndGetKeyArguments(command, key);
            classIdParameter = "key1";
            objIdParameter = "key2";
            objSubId = "2";
        }

        return $"(l.classid = @{classIdParameter} AND l.objid = @{objIdParameter} AND l.objsubid = {objSubId})";
    }

    private readonly struct CapturedTimeoutSettings(string statementTimeout, string lockTimeout)
    {
        public int StatementTimeout { get; } = ParsePostgresTimeout(statementTimeout);

        public int LockTimeout { get; } = ParsePostgresTimeout(lockTimeout);

        private static int ParsePostgresTimeout(string timeout) =>
            Regex.Match(timeout, @"^\d+(?=(?:ms)?$)") is { Success: true, Value: var value }
                ? int.Parse(value)
                : throw new FormatException($"Unexpected timeout setting value '{timeout}'");
    }
}
